package crypto

import (
    "crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"math/big"
	"strconv"

	"github.com/consensys/gnark-crypto/ecc/bls12-381"
	"github.com/consensys/gnark-crypto/ecc/bls12-381/fr"
)

// sha256HashBytes computes SHA256 of data (renamed to avoid name clash with stdlib hash pkg).
func sha256HashBytes(data []byte) []byte {
	h := sha256.New()
	h.Write(data)
	return h.Sum(nil)
}

type G1 struct{ bls12381.G1Affine }
type G2 struct{ bls12381.G2Affine }
type ZR struct{ fr.Element }
// GT is the target group of the pairing
type GT struct{ bls12381.GT }

var g1Gen G1
var g2Gen G2

func init() {
	_, _, g1Aff, g2Aff := bls12381.Generators()
	g1Gen = G1{g1Aff}
	g2Gen = G2{g2Aff}
}

func GetG1() G1 { return g1Gen }
func GetG2() G2 { return g2Gen }
func MulG1(p *G1, s *ZR) G1 {
	var resJac bls12381.G1Jac
	resJac.FromAffine(&p.G1Affine)
	resJac.ScalarMultiplication(&resJac, s.Element.BigInt(new(big.Int)))
	var resAff bls12381.G1Affine
	resAff.FromJacobian(&resJac)
	return G1{resAff}
}
func AddG1(p1, p2 *G1) G1 {
	var p1Jac, p2Jac bls12381.G1Jac
	p1Jac.FromAffine(&p1.G1Affine)
	p2Jac.FromAffine(&p2.G1Affine)
	p1Jac.AddAssign(&p2Jac)
	var resAff bls12381.G1Affine
	resAff.FromJacobian(&p1Jac)
	return G1{resAff}
}

// MulG2 multiplies G2 point by scalar s
func MulG2(p *G2, s *ZR) G2 {
    var resJac bls12381.G2Jac
    resJac.FromAffine(&p.G2Affine)
    resJac.ScalarMultiplication(&resJac, s.Element.BigInt(new(big.Int)))
    var resAff bls12381.G2Affine
    resAff.FromJacobian(&resJac)
    return G2{resAff}
}

// AddG2 adds two G2 points
func AddG2(p1, p2 *G2) G2 {
    var p1Jac, p2Jac bls12381.G2Jac
    p1Jac.FromAffine(&p1.G2Affine)
    p2Jac.FromAffine(&p2.G2Affine)
    p1Jac.AddAssign(&p2Jac)
    var resAff bls12381.G2Affine
    resAff.FromJacobian(&p1Jac)
    return G2{resAff}
}
func (z *ZR) MarshalJSON() ([]byte, error) {
	bytes := z.Element.Bytes()
	encoded := base64.StdEncoding.EncodeToString(bytes[:])
	return json.Marshal(encoded)
}
func (z *ZR) UnmarshalJSON(data []byte) error {
	var encoded string
	if err := json.Unmarshal(data, &encoded); err != nil {
		return err
	}
	bytes, err := base64.StdEncoding.DecodeString(encoded)
	if err != nil {
		return err
	}
	z.Element.SetBytes(bytes)
	return nil
}
func (g *G1) MarshalJSON() ([]byte, error) {
	bytes := g.G1Affine.Bytes()
	encoded := base64.StdEncoding.EncodeToString(bytes[:])
	return json.Marshal(encoded)
}
func (g *G1) UnmarshalJSON(data []byte) error {
	var encoded string
	if err := json.Unmarshal(data, &encoded); err != nil {
		return err
	}
	bytes, err := base64.StdEncoding.DecodeString(encoded)
	if err != nil {
		return err
	}
	if _, err := g.G1Affine.SetBytes(bytes); err != nil {
		return fmt.Errorf("failed to set G1 bytes: %w", err)
	}
	return nil
}
func (g *G2) MarshalJSON() ([]byte, error) {
	bytes := g.G2Affine.Bytes()
	encoded := base64.StdEncoding.EncodeToString(bytes[:])
	return json.Marshal(encoded)
}
func (g *G2) UnmarshalJSON(data []byte) error {
	var encoded string
	if err := json.Unmarshal(data, &encoded); err != nil {
		return err
	}
	bytes, err := base64.StdEncoding.DecodeString(encoded)
	if err != nil {
		return err
	}
	if _, err := g.G2Affine.SetBytes(bytes); err != nil {
		return fmt.Errorf("failed to set G2 bytes: %w", err)
	}
	return nil
}

// Pair computes the optimal ate pairing e(p, q)
func Pair(p *G1, q *G2) GT {
    gt, _ := bls12381.Pair([]bls12381.G1Affine{p.G1Affine}, []bls12381.G2Affine{q.G2Affine})
    return GT{gt}
}

// Equal reports whether two GT elements are equal
func (t *GT) Equal(other *GT) bool {
    return t.GT.Equal(&other.GT)
}

// SetOne sets GT to the multiplicative identity
func (t *GT) SetOne() {
    t.GT.SetOne()
}
func ZRFromInt(v int64) ZR             { var zrEl fr.Element; zrEl.SetInt64(v); return ZR{zrEl} }
func (z *ZR) Exp(a ZR, e *big.Int) *ZR { z.Element.Exp(a.Element, e); return z }
func InnerProduct(a, b []ZR) (ZR, error) {
	if len(a) != len(b) {
		return ZR{}, fmt.Errorf("vectors must have the same length")
	}
	var result ZR
	result.SetZero()
	for i := range a {
		var term ZR
		term.Mul(&a[i], &b[i])
		result.Add(&result, &term)
	}
	return result, nil
}
func EvalPolynomial(coeffs []ZR, x int64) ZR {
	var result ZR
	result.SetZero()
	xZr := ZRFromInt(x)
	for i := len(coeffs) - 1; i >= 0; i-- {
		result.Mul(&result, &xZr)
		result.Add(&result, &coeffs[i])
	}
	return result
}
func InterpolateAtX(coords [][2]interface{}, x int64) (ZR, error) {
	if len(coords) == 0 {
		return ZR{}, fmt.Errorf("no coordinates provided")
	}

	var result ZR
	result.SetZero()

	for i, coord := range coords {
		xi, ok1 := coord[0].(int64)
		yi, ok2 := coord[1].(ZR)
		if !ok1 || !ok2 {
			return ZR{}, fmt.Errorf("invalid coordinate format at index %d", i)
		}

		// Compute Lagrange basis polynomial
		var numerator ZR
		numerator.SetOne()
		var denominator ZR
		denominator.SetOne()

		for j, otherCoord := range coords {
			if i == j {
				continue
			}
			otherXi, _ := otherCoord[0].(int64)
			_, _ = otherCoord[1].(ZR) // otherYi not used in this context

			// numerator *= (x - otherXi)
			var term ZR
			term.SetInt64(x - otherXi)
			numerator.Mul(&numerator, &term)

			// denominator *= (xi - otherXi)
			var denomTerm ZR
			denomTerm.SetInt64(xi - otherXi)
			denominator.Mul(&denominator, &denomTerm)
		}

		// Compute lagrange coefficient: numerator / denominator
		var lagrangeCoeff ZR
		lagrangeCoeff.Div(&numerator, &denominator)

		// Add yi * lagrangeCoeff to result
		var term ZR
		term.Mul(&yi, &lagrangeCoeff)
		result.Add(&result, &term)
	}

	return result, nil
}

func RandomPolynomial(f int, secret ZR) []ZR {
	coeffs := make([]ZR, f+1)
	coeffs[0] = secret
	for i := 1; i <= f; i++ {
		coeffs[i].SetRandom()
	}
	return coeffs
}

func HashToG1(data []byte) G1 {
    seed := new(big.Int).SetBytes(sha256HashBytes(data))
	var secret ZR
	secret.SetBigInt(seed)
	g := GetG1()
	return MulG1(&g, &secret)
}

// HashToG2 hashes arbitrary data to a G2 point deterministically (simplified)
func HashToG2(data []byte) G2 {
    seed := new(big.Int).SetBytes(sha256HashBytes(data))
    var secret ZR
    secret.SetBigInt(seed)
    g := GetG2()
    return MulG2(&g, &secret)
}

// SetString sets a ZR from a string representation
func (z *ZR) SetString(s string) error {
	// Try to parse as integer first
	if val, err := strconv.ParseInt(s, 10, 64); err == nil {
		z.SetInt64(val)
		return nil
	}
	// Try to parse as hex
	if val, err := strconv.ParseInt(s, 16, 64); err == nil {
		z.SetInt64(val)
		return nil
	}
	return fmt.Errorf("cannot parse string %s as ZR", s)
}

// String returns a string representation of ZR
func (z *ZR) String() string {
	return z.Element.String()
}

// SetOne sets the ZR to 1
func (z *ZR) SetOne() {
	z.Element.SetOne()
}

// SetZero sets the ZR to 0
func (z *ZR) SetZero() {
	z.Element.SetZero()
}

// SetRandom sets the ZR to a random value
func (z *ZR) SetRandom() {
	z.Element.SetRandom()
}

// SetBigInt sets the ZR from a big.Int
func (z *ZR) SetBigInt(i *big.Int) {
	z.Element.SetBigInt(i)
}

// SetInt64 sets the ZR from an int64
func (z *ZR) SetInt64(i int64) {
	z.Element.SetInt64(i)
}

// Add adds two ZR elements
func (z *ZR) Add(a, b *ZR) *ZR {
	z.Element.Add(&a.Element, &b.Element)
	return z
}

// Sub subtracts two ZR elements
func (z *ZR) Sub(a, b *ZR) *ZR {
	z.Element.Sub(&a.Element, &b.Element)
	return z
}

// Mul multiplies two ZR elements
func (z *ZR) Mul(a, b *ZR) *ZR {
	z.Element.Mul(&a.Element, &b.Element)
	return z
}

// Div divides two ZR elements
func (z *ZR) Div(a, b *ZR) *ZR {
	z.Element.Div(&a.Element, &b.Element)
	return z
}

// Neg negates the ZR element
func (z *ZR) Neg(a *ZR) *ZR {
	z.Element.Neg(&a.Element)
	return z
}

// Inv computes the multiplicative inverse
func (z *ZR) Inv(a *ZR) *ZR {
	z.Element.Inverse(&a.Element)
	return z
}

// Equal checks if two ZR elements are equal
func (z *ZR) Equal(other *ZR) bool {
	return z.Element.Equal(&other.Element)
}

// IsZero checks if the ZR element is zero
func (z *ZR) IsZero() bool {
	return z.Element.IsZero()
}

// IsOne checks if the ZR element is one
func (z *ZR) IsOne() bool {
	return z.Element.IsOne()
}
